import os
import json

import cv2
from ultralytics import YOLO
from PIL import Image
import extract_seal as es
import seal_recogniton as sr


def process_yolo_predictions(input_folder, model_path,extract_folder="cropped_seals", save_dict_json=False):
    """
    处理YOLO预测结果并按要求输出

    参数:
        input_folder: 包含输入图片的文件夹路径
        model_path: 训练好的YOLO模型路径(.pt文件)
        save_dict_json: 是否将类别字典保存为JSON文件
    """
    # 加载模型
    model = YOLO(model_path)

    # 类别映射
    class_mapping = {
        0: "red_off_seal",
        1: "red_person_seal",
        2: "black_off_seal",
        3: "black_person_seal"
    }

    # 创建输出文件夹
    os.makedirs(extract_folder, exist_ok=True)

    # 初始化结果字典
    result_dict = {}

    # 遍历输入文件夹中的所有图片
    for img_name in os.listdir(input_folder):
        if not img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
            continue

        img_path = os.path.join(input_folder, img_name)
        results = model.predict(img_path)

        # 获取基本文件名(不带扩展名)
        base_name = os.path.splitext(img_name)[0]

        # 存储该图片中检测到的类别
        detected_classes = set()

        # 处理每个检测结果
        for i, result in enumerate(results):
            boxes = result.boxes
            orig_img = result.orig_img
            image_list =[]

            for j, box in enumerate(boxes):
                # 获取类别ID
                cls_id = int(box.cls)
                detected_classes.add(cls_id)

                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, box.xyxy[0])

                # 裁剪图像
                cropped_img = orig_img[y1:y2, x1:x2]

                # 掩码提取
                cropped_img=es.extract_seal_array(image=cropped_img, h_scale=[150, 30], s_scale=[20, 255], v_scale=[20, 255])

                # 确保颜色通道正确 - 转换为RGB
                if cropped_img.shape[-1] == 3:  # 如果是彩色图像
                    cropped_img = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)

                # 生成保存名称
                save_name = f"{base_name}-{class_mapping[cls_id]}-{cls_id}-{j}.jpg"
                image_list.append(save_name);
                save_path = os.path.join(extract_folder, save_name)

                # 保存裁剪图像
                Image.fromarray(cropped_img).save(save_path)

            # # 全部剪裁文件生成后， 执行seal_recogniton
            # print(base_name+" 裁剪文件生成完毕")
            # sr.process_seal_recognition_bylist(input_img_list=image_list, output_dir="output")


        # 将检测到的类别添加到结果字典
        result_dict[img_name] = list(detected_classes)

    # 如果需要保存为JSON
    if save_dict_json:
        with open("detection_results.json", "w") as f:
            json.dump(result_dict, f, indent=4)

    print("yolo检测 处理完成")

    return result_dict

# 使用示例
# result = process_yolo_predictions("seal_examples", "seal_sign_model.pt", save_dict_json=True)
